"""
Answe inspired by the post from here: 
https://leetcode.com/problems/house-robber-iii/discuss/79330/Step-by-step-tackling-of-the-problem
"""

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class Solution:
    # Method 2: Greedy approach:
    def rob2(self, root):
        def helper(root):
            if not root: return [0, 0]
            left, right = helper(root.left), helper(root.right)
            not_robbed = max(left) + max(right)
            robbed = root.val + left[0] + right[0]
            return [not_robbed, robbed]
        return max(helper(root))

case1 = TreeNode(3, 
                 left = TreeNode(2, None, TreeNode(3)),
                 right = TreeNode(3, None, TreeNode(1)))
case2 = TreeNode(3,
                 left = TreeNode(4, TreeNode(1), TreeNode(3)),
                 right = TreeNode(5, None, TreeNode(1)))
s = Solution()
print(s.rob2(case1))
print(s.rob2(case2))